﻿#include "im3d_conversion.h"
#include "im3d_render.h"
#include <numeric>
#include <cmath>
#include <format>

IM3DConversion::IM3DConversion(const std::string& model_path, bool use_bf16,
    int width, int height, float x_shift)
    : model_path(model_path), use_bf16_(use_bf16), frame_width_(width), frame_height_(height),
    x_shift_(x_shift* width), base_grid_(frame_width_, frame_height_){

    // 初始化视差序列
    int num_planes = 16;
    disparities_.resize(num_planes);
    float step = (1.0f - 0.001f) / (num_planes + 1);
    for (int i = 0; i < num_planes; ++i) {
        disparities_[i] = 1.0f - step * (i + 1);
    }
}

IM3DConversion::~IM3DConversion() {}

bool IM3DConversion::initialize() {
    try {
        // 加载模型
        std::string model_name = use_bf16_ ? "lmpi2_bf16_traced.xml" : "lmpi2_traced.xml";
        std::string model_filename = model_path + '/' + model_name;
        model_ = core_.read_model(model_filename);

        // 配置输入输出
        auto inputs = model_->inputs();
        auto outputs = model_->outputs();
        input_name_ = inputs[0].get_any_name();
        output_name_ = outputs[0].get_any_name();

        // 编译模型
        compiled_model_ = core_.compile_model(model_, "CPU");
        infer_request_ = compiled_model_.create_infer_request();

        // 创建基础网格和预计算网格
        //createBaseGrid();
        CoordinateGrid grid(frame_width_, frame_height_);
        grid.print_shape();
        //precomputeSamplingGrid();

        std::cout << "IM3DConversion initialized successfully" << std::endl;
        return true;

    }
    catch (const std::exception& e) {
        std::cerr << "Initialization failed: " << e.what() << std::endl;
        return false;
    }
}

/*
void IM3DConversion::precomputeSamplingGrid() {
    int num_planes = disparities_.size();
    flat_grid_precomputed_ = cv::Mat(num_planes * frame_height_, frame_width_, CV_32FC2);

    std::cout << "flat_grid_precomputed size" << flat_grid_precomputed_.size() << std::endl;

    for (int s = 0; s < num_planes; ++s) {
        float plane_shift = x_shift_ * disparities_[s];
        float x_offset = plane_shift * (2.0f / frame_width_);

        for (int y = 0; y < frame_height_; ++y) {
            for (int x = 0; x < frame_width_; ++x) {
                std::pair<float, float> base_coord = base_grid_(y, x);
                cv::Vec2f offset_coord(base_coord.first - x_offset, base_coord.second);
                flat_grid_precomputed_.at<cv::Vec2f>(s * frame_height_ + y, x) = offset_coord;
            }
        }
    }
}
*/

cv::Mat IM3DConversion::processFrame(const cv::Mat& frame) {
    try {
        // 输入图像预处理
        cv::Mat input_frame;
        //cv::resize(frame, input_frame, cv::Size(frame_width_, frame_height_));
        int org_height, org_width, channels_;
        org_height = frame.rows;
        org_width = frame.cols;
        channels_ = frame.channels();
        input_frame = frame.clone();

        // 水平方向采样一半
        std::cout << "half crop frame " << std::endl;
        int half_width = org_width / 2;
        cv::Mat image_half_width;
        cv::resize(input_frame, image_half_width, cv::Size(half_width, frame_height_));

        // 转换为模型输入尺寸
        cv::Mat image_low_res;
        cv::resize(image_half_width, image_low_res, cv::Size(mpi_width_, mpi_height_));

        // 转换为浮点并归一化
        cv::Mat float_image;
        image_low_res.convertTo(float_image, CV_32FC3, 1.0 / 255.0);

        // 调整通道顺序 HWC -> CHW
        std::vector<cv::Mat> channels(3);
        cv::split(float_image, channels);
        std::cout << "Change HWC to CHW " << std::endl;

        // 创建CHW格式的图像
        cv::Mat chw_image(mpi_height_ * 3, mpi_width_, CV_32FC1);
        for (int i = 0; i < 3; ++i) {
            channels[i].reshape(1, mpi_height_).copyTo(chw_image.rowRange(i * mpi_height_, (i + 1) * mpi_height_));
        }

       
        // 创建输入张量
        ov::Shape input_shape = { 1, 3, (unsigned long long)mpi_height_,  (unsigned long long)mpi_width_ };
        std::cout << "ov input shape " << input_shape << std::endl;
        ov::Tensor input_tensor(ov::element::f32, input_shape);

        // 复制数据到输入张量
        float* input_data = input_tensor.data<float>();
        memcpy(input_data, chw_image.data, mpi_height_ * mpi_width_ * 3 * sizeof(float));

        std::cout << "before infer " << std::endl;
        // 推理
        infer_request_.set_input_tensor(input_tensor);
        infer_request_.infer();
        std::cout << "infer success " << std::endl;

        // 获取输出
        ov::Tensor output_tensor = infer_request_.get_output_tensor();
        const float* output_data = output_tensor.data<const float>();
        ov::Shape shape = output_tensor.get_shape();
        std::cout << "output shape: [";
        for (size_t i = 0; i < shape.size(); ++i) {
            std::cout << shape[i];
            if (i < shape.size() - 1) std::cout << ", ";
        }
        std::cout << "]" << std::endl;
        std::cout << "get output data size " << output_data << std::endl;

        // 解析MPI输出 [16, 4, mpi_height, mpi_width]
        std::vector<cv::Mat> mpi_planes;

        // 检查输出形状是否正确
        if (shape.size() != 5) {
            throw std::runtime_error("Unexpected output shape from model. Expected 5D tensor [1,16,4,H,W]");
        }

        int batch_size = shape[0];     // 1
        int num_planes = shape[1];     // 16
        int num_channels = shape[2];   // 4
        int mpi_h = shape[3];          // 256
        int mpi_w = shape[4];          // 384

        std::cout << "MPI output: batch=" << batch_size << ", " << num_planes << " planes, "
            << num_channels << " channels, " << mpi_h << " x " << mpi_w << std::endl;

        if (batch_size != 1) {
            throw std::runtime_error("Only batch size 1 is supported");
        }

        // 计算正确的数据偏移量
        size_t batch_stride = num_planes * num_channels * mpi_h * mpi_w;  // 整个批次的大小
        size_t plane_stride = num_channels * mpi_h * mpi_w;               // 单个平面的大小
        size_t channel_stride = mpi_h * mpi_w;                            // 单个通道的大小

        // 跳过批次维度，直接访问第一个批次的数据
        const float* batch_data = output_data;


        int plane_size = 4 * mpi_height_ * mpi_width_;

        for (int plane_idx = 0; plane_idx < num_planes; ++plane_idx) {
            // 为每个平面创建4通道图像
            cv::Mat plane(mpi_h, mpi_w, CV_32FC4);
            float* plane_data = reinterpret_cast<float*>(plane.data);

            // 复制该平面的所有4个通道数据
            const float* src_data = batch_data + plane_idx * plane_stride;

            // 逐通道复制，确保正确的内存布局
            for (int ch = 0; ch < num_channels; ++ch) {
                const float* channel_src = src_data + ch * channel_stride;
                float* channel_dst = plane_data + ch;

                // 复制整个通道数据（交错存储）
                for (int i = 0; i < mpi_h * mpi_w; ++i) {
                    channel_dst[i * num_channels] = channel_src[i];
                }
            }

            mpi_planes.push_back(plane);

            // 调试：保存每个平面的可视化
            //debugSaveMPIPlane(plane, plane_idx);
        }

        std::cout << "get mpi plane " << mpi_planes.size() << std::endl;

        // 上采样MPI到目标尺寸
        std::vector<cv::Mat> upsampled_mpi;
        int plane_idx = 0;
        for (auto& plane : mpi_planes) {
            cv::Mat upsampled;
            cv::resize(plane, upsampled, cv::Size(half_width, frame_height_));
            upsampled_mpi.push_back(upsampled);
            plane_idx++;
        }
        std::cout << "upsample mpi finished start render frame" << std::endl;

     
        cv::Mat base_grid = IM3DRenderer::createBaseGrid(half_width, frame_height_);
        std::vector<float> disparities = IM3DRenderer::generateDisparities(16, 1.0f, 0.001f);
        float actual_x_shift = 0.014f * half_width; // 注意：这里使用half_width而不是frame_width_
        Grid4D grid_4d = IM3DRenderer::precomputeSamplingGrid4D(disparities, base_grid, actual_x_shift);

        std::cout << "Created 4D grid with shape: [" << grid_4d.getNumPlanes() << ", "
            << grid_4d.getHeight() << ", " << grid_4d.getWidth() << ", 2]" << std::endl;

        //std::cout << "Created sampling grid with size: " << flat_grid.size() << std::endl;
        std::cout << "Upsampled MPI planes: " << upsampled_mpi.size() << std::endl;
        std::cout << "First MPI plane size: " << upsampled_mpi[0].size() << std::endl;
        cv::Mat r_tgt = IM3DRenderer::renderFrame(upsampled_mpi, image_half_width, grid_4d);
        std::cout << "render frame finished" << std::endl;

        // 拼接并排图像
        std::cout << "r_tgt size " << r_tgt.size() << "r_tgt type " << r_tgt.type() << "r_tgt channel" << r_tgt.channels() 
            << "image_half_width size " << image_half_width.size() << " type" << image_half_width.type() << "channel " << image_half_width.channels() << std::endl;

        if (r_tgt.type() != image_half_width.type()) {
            std::cerr << "Error: Data type mismatch! r_tgt.type=" << r_tgt.type()
                << ", image_half_width.type=" << image_half_width.type() << std::endl;

            // 转换数据类型使其匹配
            if (r_tgt.type() == CV_32FC3) {
                image_half_width.convertTo(image_half_width, CV_32FC3, 1.0 / 255.0);
            }
            else if (r_tgt.type() == CV_8UC3) {
                r_tgt.convertTo(r_tgt, CV_8UC3);
            }
            std::cout << "Converted data types to match" << std::endl;
        }

        cv::Mat sbs_frame;
        cv::hconcat(std::vector<cv::Mat>{r_tgt, image_half_width}, sbs_frame);
        sbs_frame = 255 * sbs_frame;
        return sbs_frame;

    }
    catch (const std::exception& e) {
        std::cerr << "Processing failed: " << e.what() << std::endl;
        return cv::Mat();
    }
}

cv::Mat IM3DConversion::renderFrame(const std::vector<cv::Mat>& mpi, const cv::Mat& org_image) {
    int num_planes = mpi.size();
    int height = mpi[0].rows;
    int width = mpi[0].cols;

    // 简化的alpha合成渲染
    cv::Mat result = cv::Mat::zeros(height, width, CV_32FC3);
    cv::Mat cumulative_transparency = cv::Mat::ones(height, width, CV_32FC1);

    // 将原始图像转换为浮点
    cv::Mat org_float;
    org_image.convertTo(org_float, CV_32FC3, 1.0 / 255.0);

    for (int i = 0; i < num_planes; ++i) {
        // 分离MPI平面的通道
        std::vector<cv::Mat> plane_channels;
        cv::split(mpi[i], plane_channels);

        // RGB通道 (前3个通道)
        cv::Mat rgb_plane;
        std::vector<cv::Mat> rgb_channels(plane_channels.begin(), plane_channels.begin() + 3);
        cv::merge(rgb_channels, rgb_plane);

        // Alpha通道 (第4个通道)
        cv::Mat alpha = plane_channels[3];

        // 使用原始图像更新RGB通道 (模拟torch.lerp)
        cv::Mat updated_rgb;
        for (int c = 0; c < 3; ++c) {
            cv::Mat updated_channel;
            // result = start + weight * (end - start)
            cv::subtract(org_float.channels() > c ? getChannel(org_float, c) : cv::Mat(),
                plane_channels[c], updated_channel);
            cv::multiply(updated_channel, alpha, updated_channel);
            cv::add(plane_channels[c], updated_channel, updated_channel);

            if (c == 0) {
                updated_rgb = updated_channel;
            }
            else {
                cv::merge(std::vector<cv::Mat>{updated_rgb, updated_channel}, updated_rgb);
            }
        }

        // 当前平面的贡献
        cv::Mat plane_contribution;
        cv::multiply(updated_rgb, alpha, plane_contribution);
        cv::multiply(plane_contribution, cumulative_transparency, plane_contribution);

        // 累加到结果
        cv::add(result, plane_contribution, result);

        // 更新累积透明度
        cv::Mat one_minus_alpha;
        cv::subtract(cv::Scalar(1.0), alpha, one_minus_alpha);
        cv::multiply(cumulative_transparency, one_minus_alpha, cumulative_transparency);
    }

    // 转换为8位
    cv::Mat final_result;
    result.convertTo(final_result, CV_8UC3, 255.0);

    return final_result;
}

// 辅助函数：安全地获取通道
cv::Mat IM3DConversion::getChannel(const cv::Mat& image, int channel) {
    std::vector<cv::Mat> channels;
    cv::split(image, channels);
    if (channel < channels.size()) {
        return channels[channel];
    }
    return cv::Mat();
}